#!/usr/bin/env python3
"""
增强版YouTube Cookies设置工具
支持多种浏览器cookie导出方法
"""

import os
import sys
import subprocess
import json
from pathlib import Path

def check_browser_cookies():
    """检查可用的浏览器cookies导出方法"""
    print("🍪 YouTube Cookies 设置指南")
    print("=" * 60)
    
    methods = [
        {
            "name": "方法1：使用yt-dlp自动提取 (推荐)",
            "description": "让yt-dlp自动从浏览器提取cookies",
            "command": "yt-dlp --cookies-from-browser chrome https://youtube.com/shorts/KdOZEnULbOA",
            "note": "支持chrome, firefox, safari, edge等浏览器"
        },
        {
            "name": "方法2：浏览器扩展导出",
            "description": "使用浏览器扩展手动导出cookies",
            "extensions": [
                "Chrome: cookies.txt (Chrome扩展)",
                "Firefox: cookies.txt (Firefox扩展)",
                "Edge: cookies.txt (Edge扩展)"
            ],
            "note": "导出后保存为 youtube_cookies.txt"
        },
        {
            "name": "方法3：手动登录测试",
            "description": "测试当前系统是否已有可用的cookies",
            "command": "yt-dlp --cookies-from-browser chrome --simulate https://youtube.com/shorts/KdOZEnULbOA",
            "note": "验证cookies是否有效"
        }
    ]
    
    for i, method in enumerate(methods, 1):
        print(f"\n{method['name']}")
        print("-" * 40)
        print(f"📝 {method['description']}")
        
        if 'command' in method:
            print(f"💻 命令: {method['command']}")
        
        if 'extensions' in method:
            print("🔧 浏览器扩展:")
            for ext in method['extensions']:
                print(f"   • {ext}")
        
        print(f"💡 {method['note']}")

def auto_extract_cookies():
    """自动提取浏览器cookies"""
    print("\n🤖 尝试自动提取浏览器cookies...")
    
    browsers = ['chrome', 'firefox', 'safari', 'edge']
    test_url = "https://youtube.com/shorts/KdOZEnULbOA"
    
    for browser in browsers:
        print(f"\n🔍 尝试从 {browser.title()} 提取cookies...")
        
        try:
            # 构造yt-dlp命令
            cmd = [
                'yt-dlp',
                '--cookies-from-browser', browser,
                '--write-info-json',
                '--skip-download',
                '--no-warnings',
                test_url
            ]
            
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
            
            if result.returncode == 0:
                print(f"✅ 成功从 {browser.title()} 提取cookies!")
                
                # 生成cookies文件命令
                cookie_cmd = [
                    'yt-dlp',
                    '--cookies-from-browser', browser,
                    '--cookies', './youtube_cookies.txt',
                    '--skip-download',
                    test_url
                ]
                
                print(f"💾 生成cookies文件...")
                cookie_result = subprocess.run(cookie_cmd, capture_output=True, text=True, timeout=30)
                
                if cookie_result.returncode == 0 and os.path.exists('./youtube_cookies.txt'):
                    print(f"🎉 成功生成 youtube_cookies.txt!")
                    return True
                else:
                    print(f"⚠️ cookies文件生成失败")
            else:
                print(f"❌ {browser.title()} 提取失败: {result.stderr.strip()}")
                
        except subprocess.TimeoutExpired:
            print(f"⏰ {browser.title()} 提取超时")
        except FileNotFoundError:
            print(f"❌ 未找到yt-dlp命令")
            return False
        except Exception as e:
            print(f"❌ {browser.title()} 提取出错: {str(e)}")
    
    return False

def manual_cookie_setup():
    """手动设置cookies指南"""
    print("\n📖 手动设置Cookies指南")
    print("=" * 60)
    
    steps = [
        "1. 在浏览器中打开 https://youtube.com 并登录",
        "2. 安装浏览器cookies导出扩展:",
        "   • Chrome: 搜索 'cookies.txt' 扩展",
        "   • Firefox: 搜索 'cookies.txt' 扩展",
        "3. 访问任意YouTube视频页面",
        "4. 使用扩展导出cookies，选择 'youtube.com' 域名",
        "5. 将导出的cookies文件保存为 'youtube_cookies.txt'",
        "6. 将文件放在项目根目录中"
    ]
    
    for step in steps:
        print(f"   {step}")
    
    print("\n💡 提示:")
    print("   • cookies文件应包含 youtube.com 域名的认证信息")
    print("   • 确保在浏览器中已经登录YouTube账号")
    print("   • 某些地区可能需要代理才能正常访问")

def test_cookies_file():
    """测试cookies文件是否有效"""
    cookies_file = './youtube_cookies.txt'
    
    if not os.path.exists(cookies_file):
        print(f"❌ 未找到cookies文件: {cookies_file}")
        return False
    
    print(f"\n🧪 测试cookies文件有效性...")
    test_url = "https://youtube.com/shorts/KdOZEnULbOA"
    
    try:
        cmd = [
            'yt-dlp',
            '--cookies', cookies_file,
            '--simulate',
            '--no-warnings',
            test_url
        ]
        
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
        
        if result.returncode == 0:
            print("✅ cookies文件有效！可以正常访问YouTube")
            return True
        else:
            print(f"❌ cookies文件测试失败: {result.stderr.strip()}")
            return False
            
    except Exception as e:
        print(f"❌ 测试过程出错: {str(e)}")
        return False

def create_test_script():
    """创建带cookies的测试脚本"""
    test_script = """#!/usr/bin/env python3
import requests
import json

def test_with_cookies():
    url = "http://localhost:8000/youtube/professional-download"
    
    test_data = {
        "youtube_url": "https://youtube.com/shorts/KdOZEnULbOA?si=2DAqfdvcrDst_G6a",
        "quality": "1080p", 
        "upload_to_qiniu": False,
        "proxy_url": "http://127.0.0.1:7890",
        "cookies_file": "./youtube_cookies.txt"
    }
    
    print("🚀 测试带cookies的YouTube下载...")
    print(f"🍪 Cookies文件: {test_data['cookies_file']}")
    
    try:
        response = requests.post(url, json=test_data, timeout=300)
        result = response.json()
        
        if result.get("code") == 200:
            data = result.get("data", {})
            print("✅ 下载成功!")
            print(f"📺 标题: {data.get('title', 'Unknown')}")
            print(f"🍪 使用cookies: {data.get('cookies_used', '无')}")
        else:
            print(f"❌ 下载失败: {result.get('msg', '未知错误')}")
            
    except Exception as e:
        print(f"💥 测试出错: {str(e)}")

if __name__ == "__main__":
    test_with_cookies()
"""
    
    with open('test_with_cookies.py', 'w', encoding='utf-8') as f:
        f.write(test_script)
    
    print("\n📄 已创建测试脚本: test_with_cookies.py")
    print("   使用命令运行: python test_with_cookies.py")

def main():
    print("🍪 YouTube Cookies 增强设置工具")
    print("解决 'Sign in to confirm you're not a bot' 问题")
    print("=" * 60)
    
    # 检查yt-dlp是否可用
    try:
        result = subprocess.run(['yt-dlp', '--version'], capture_output=True, text=True)
        print(f"✅ yt-dlp 版本: {result.stdout.strip()}")
    except FileNotFoundError:
        print("❌ 未找到yt-dlp，请先安装: pip install yt-dlp")
        return
    
    # 显示设置方法
    check_browser_cookies()
    
    # 询问用户选择
    print("\n" + "=" * 60)
    choice = input("选择操作 (1:自动提取 2:手动指南 3:测试cookies 4:全部): ").strip()
    
    if choice == '1' or choice == '4':
        success = auto_extract_cookies()
        if success:
            test_cookies_file()
    
    if choice == '2' or choice == '4':
        manual_cookie_setup()
    
    if choice == '3' or choice == '4':
        test_cookies_file()
    
    # 创建测试脚本
    create_test_script()
    
    print("\n" + "=" * 60)
    print("🎯 总结:")
    print("1. 确保有效的 youtube_cookies.txt 文件")
    print("2. 使用代理（如需要）: http://127.0.0.1:7890")
    print("3. 在API请求中包含 cookies_file 参数")
    print("4. 运行 test_with_cookies.py 验证功能")

if __name__ == "__main__":
    main()